package com.example.demo.Service;

import com.volcengine.ark.runtime.model.bot.completion.chat.BotChatCompletionRequest;
import com.volcengine.ark.runtime.model.bot.completion.chat.BotChatCompletionResult;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class ArkServiceClient {

    @Value("${ark.api.key}")
    private String apiKey;

    @Value("${ark.api.url}")
    private String baseUrl;

    // 用来存储不同用户的会话历史
    private Map<String, List<ChatMessage>> userConversations = new ConcurrentHashMap<>();

    public String getBotResponse(String botId, String userId, String userMessage) {
        // 如果该用户的对话历史不存在，则初始化一个新的会话历史
        userConversations.putIfAbsent(userId, new ArrayList<>());

        // 获取该用户的对话历史
        List<ChatMessage> messages = userConversations.get(userId);

        // 添加系统消息和用户消息到历史中
        messages.add(ChatMessage.builder().role(ChatMessageRole.USER).content(userMessage).build());

        // 创建 ArkService 客户端
        ArkService service = ArkService.builder()
                .apiKey(apiKey)
                .baseUrl(baseUrl)
                .build();

        // 创建聊天请求，将所有历史消息传递给 API
        BotChatCompletionRequest chatCompletionRequest = BotChatCompletionRequest.builder()
                .botId(botId)
                .messages(messages)
                .build();

        // 获取聊天响应
        BotChatCompletionResult chatCompletionResult = service.createBotChatCompletion(chatCompletionRequest);

        // 获取机器人的回复，并将其添加到对话历史中
        StringBuilder response = new StringBuilder();
        chatCompletionResult.getChoices().forEach(choice -> {
            String botReply = (String) choice.getMessage().getContent();
            response.append(botReply).append("\n");

            // 保存机器人的回答到对话历史
            messages.add(ChatMessage.builder().role(ChatMessageRole.ASSISTANT).content(botReply).build());
        });

        // 关闭服务
        service.shutdownExecutor();

        // 返回机器人的回应
        return response.toString();
    }
}
